import pandas as pd
import os
from joblib import load
from flask import jsonify

class Processor:
    def __init__(self, model_path, scaler_path, imputer_path):
        self.model = load(model_path)
        self.scaler = load(scaler_path)
        self.imputer = load(imputer_path)

    def preprocess_data(self, df):
        # 确保列名与训练时完全一致（去除空格）
        df.columns = df.columns.str.strip()
        
        # 特征选择
        required_features = ['红细胞计数', '血红蛋白浓度', '红细胞压积', '平均红细胞体积', 
                            '平均红细胞血红蛋白含量', '平均红细胞血红蛋白浓度', '红细胞体积分布宽度']
        X = df[required_features]
        
        # 处理缺失值（使用训练时的Imputer）
        if X.isnull().sum().sum() > 0:
            X = self.imputer.transform(X)
        else:
            X = X.values  # 转换为NumPy数组
        
        # 确保特征名称与训练时一致
        if hasattr(self.scaler, 'feature_names_in_'):
            if len(X.shape) == 1:
                X = X.reshape(1, -1)
            X = pd.DataFrame(X, columns=self.scaler.feature_names_in_)
        
        # 标准化（使用训练时的Scaler）
        X = self.scaler.transform(X)
        
        return X

    def save_to_csv(self, patientId, recordId, data, predictions, probabilities):
        # 确保temp目录存在
        patient_data_dir = 'temp/{patientId}'.format(patientId=patientId)
        if not os.path.exists(patient_data_dir):
            os.makedirs(patient_data_dir)
        
        # 构造文件路径
        file_path = os.path.join(patient_data_dir, f'{recordId}.csv')
        
        # 构造要保存的数据
        df = pd.DataFrame(data)
        df['recordId'] = recordId
        df['预测'] = predictions
        df['地中海贫血概率'] = [prob[0] for prob in probabilities]
        df['正常概率'] = [prob[1] for prob in probabilities]
        df['缺铁性贫血概率'] = [prob[2] for prob in probabilities]
        
        # 确保列顺序正确
        columns = ['recordId', '红细胞计数', '血红蛋白浓度', '红细胞压积', '平均红细胞体积',
                   '平均红细胞血红蛋白含量', '平均红细胞血红蛋白浓度', '红细胞体积分布宽度',
                   '预测', '地中海贫血概率', '正常概率', '缺铁性贫血概率']
        df = df[columns]
        
        # 保存到CSV文件
        if os.path.exists(file_path):
            # 如果文件存在，追加数据
            df.to_csv(file_path, mode='a', header=False, index=False, encoding='utf-8-sig')
        else:
            # 如果文件不存在，创建并写入表头
            df.to_csv(file_path, index=False, encoding='utf-8-sig')

    def predict_and_save(self, patientId, recordId, bloodRoutineIndicators):
        df = pd.DataFrame(bloodRoutineIndicators)
        
        # 预处理数据
        try:
            X_processed = self.preprocess_data(df)
        except KeyError as e:
            return jsonify({"error": f"缺失必要特征列: {str(e)}"}), 400
        
        # 预测
        predictions = self.model.predict(X_processed)
        probabilities = self.model.predict_proba(X_processed)
        
        # 保存到CSV文件
        self.save_to_csv(patientId, recordId, bloodRoutineIndicators, predictions, probabilities)
        
        result = {
            'recordId': recordId,
            'predictions': predictions.tolist(),
            'probabilities': probabilities.tolist()
        }
        
        return jsonify(result)